'''
Created on Nov 1, 2011

@author: alebalbin
'''
import numpy as np
from collections import OrderedDict, defaultdict

def read_ids(nfile):
    ifile=open(nfile)
    h=ifile.next().strip('\r\n').split('\t')
    gids=defaultdict(list)
    i=0
    for l in ifile:
        f=l.strip('\r\n').split('\t')
        gids[f[0]].append(i)
        i+=1
    ifile.close()
    return gids,h

def calc_fold_change(mat,offset, cind):
    '''
    group2=treament, group1=control, group2/group1
    '''
    fold_change = np.true_divide( np.mean(mat[:,(cind+1):]+offset,axis=1), np.mean(mat[:,:cind]+offset,axis=1) )  
    return fold_change
    
def find_repeated_features(features):
    '''
    '''
    dind=np.array(range(len(features)))
    d=np.array([features.count(i) for i in features])
    return list(d[d>1])
    
def summaryze_repeat_features(mat, gids):
    '''
    summaryze repeated features using the features with highest sum across samples
    '''
    nmat = np.empty( (len(gids), mat.shape[1]) )
    ngids=defaultdict(list)
    i=0
    for f,rows in gids.iteritems():
        rows=np.array(rows)
        s = np.sum(mat[rows,],axis=1)
        s = rows[s==np.max(s)]
        if len(s)>1:
            s=s[0]
        #nmat[i,:]=np.mean(mat[rows,],axis=0) tp get the mean summary
        nmat[i,:]=mat[s,:]
        ngids[f].append(i)
        i+=1
    return nmat, ngids

    
def read_mat(nfile,fcol, cind,offset):
    '''
    '''
    gids,h=read_ids(nfile)
    mat = np.loadtxt(nfile,skiprows=1,usecols=set(range(fcol,len(h) )))
    mat,gids = summaryze_repeat_features(mat, gids)
    #
    fd=calc_fold_change(mat,offset, cind)
    mat = np.append(mat,np.reshape(fd,(mat.shape[0],-1)),axis=1)
    
    return mat, gids

def intersect_ids(gids_list):
    '''
    '''
    isec_gids=set(gids_list[0])
    for gids in gids_list:
        isec_gids.intersection_update(gids)
    return isec_gids
    
    
def intersect_matrices(file1,file2,file3,ofile,dolog2=True):    
    '''
    '''
    fcol, offset, nochange=2,0.5,-999
    pind=5
    tind=6
    trans_mat,tgids = read_mat(file1,fcol, tind,offset)
    prots_mat, pgids = read_mat(file2,fcol, pind,offset)
    phospho_mat, phosgids = read_mat(file3,fcol, pind,offset)
    # isec gids
    
    # First: common to all data sets
    all_ids = [tgids.keys(),pgids.keys(),phosgids.keys()]
    isec_gids = intersect_ids(all_ids)
    print len(isec_gids),len( set(sum(all_ids,[] ))), len(tgids.keys())
    fccol=-1 # fold change column, is the last one for default
    
    cmat = np.empty( (len(set(sum(all_ids,[]))), len(all_ids) ) ) 
    
    of=open(ofile,'w')
    hd=["NAME","DESCRIPTION","transcriptFoldChange","proteinFoldChange","phosphoFoldChange"]
    of.write(",".join(hd).replace(',', '\t')+'\n')
    i=0
    for f in isec_gids:
        
        t,p,pp=tgids[f],pgids[f],phosgids[f]
        #store only the fold change value. The same can be done for the whole matrix
        cmat[i,0],cmat[i,1],cmat[i,2]=trans_mat[t,fccol],prots_mat[p,fccol],phospho_mat[pp,fccol]
        if dolog2:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,np.log2(cmat[i,:])))).replace(',','\t')+'\n')
        else:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,cmat[i,:]))).replace(',','\t')+'\n')
        i+=1
    
    # Second: common to transcript and phospho
    tpall_ids = [tgids.keys(),phosgids.keys()]
    tphos_ids = intersect_ids(tpall_ids)
    tphos_ids.difference_update(isec_gids)
    print len(tphos_ids)
    
    fccol=-1 # fold change column, is the last one for default
    for f in tphos_ids:
        t,pp=tgids[f],phosgids[f]
        #store only the fold change value. The same can be done for the whole matrix
        cmat[i,0],cmat[i,1],cmat[i,2]=trans_mat[t,fccol],nochange,phospho_mat[pp,fccol]
        if dolog2:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,np.log2(cmat[i,:])))).replace(',','\t')+'\n')
        else:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,cmat[i,:]))).replace(',','\t')+'\n')
            
        i+=1
    
    # Third: common to transcript and protein 
    
    tpall_ids = [tgids.keys(),pgids.keys()]
    tprot_ids = intersect_ids(tpall_ids)
    tprot_ids.difference_update(isec_gids)
    print len(tprot_ids)

    
    fccol=-1 # fold change column, is the last one for default
    for f in tprot_ids:
        
        t,p=tgids[f],pgids[f]
        #store only the fold change value. The same can be done for the whole matrix
        cmat[i,0],cmat[i,1],cmat[i,2]=trans_mat[t,fccol],prots_mat[p,fccol],nochange
        if dolog2:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,np.log2(cmat[i,:])))).replace(',','\t')+'\n')
        else:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,cmat[i,:]))).replace(',','\t')+'\n')
            
        i+=1
    
    tgids_uq = set(tgids.keys()).difference(isec_gids)
    tgids_uq = tgids_uq.difference(tphos_ids)
    tgids_uq = tgids_uq.difference(tprot_ids)
    
    for f in tgids_uq:
        t=tgids[f]
        #store only the fold change value. The same can be done for the whole matrix
        cmat[i,0],cmat[i,1],cmat[i,2]=trans_mat[t,fccol],nochange,nochange
        if dolog2:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,np.log2(cmat[i,:])))).replace(',','\t')+'\n')
        else:
            of.write(f+'\t'+"NAN"+'\t'+",".join(list(map(str,cmat[i,:]))).replace(',','\t')+'\n')
            
        i+=1

        
    of.close()    
    
if __name__ == '__main__':
    
    file1='/Users/alebalbin/Desktop/nsclc_kras_dep/Data/transcriptomics/2010_07_10_21_54_lung_kras_allcurated_log_notisof_outfile_expression_matrix_pamr_onlyRas_USEFOR_GSEA.txt'
    file2='/Users/alebalbin/Desktop/nsclc_kras_dep/Data/phosphoproteomics/phosphoproteomicsdatasets/H2122_dep/flowt_norm_prots2_gsea.txt'
    file3='/Users/alebalbin/Desktop/nsclc_kras_dep/Data/phosphoproteomics/phosphoproteomicsdatasets/H2122_dep/phospho_norm_prots2_gsea.txt'
    ofile='/Users/alebalbin/Desktop/nsclc_kras_dep/Data/phosphoproteomics/phosphoproteomicsdatasets/H2122_dep/full_phospho_flow_trans_KDEP_to_KIND_FC_merged_mat.tsv'
    dolog2=False
    # Missing measurements at any of the levels Transcript, Protein, PhosphoProt are indicated by a perfect -999
    intersect_matrices(file1,file2,file3,ofile,dolog2)

    
    
    